import torch
from thop import profile
from models.BiFNet.BiFNet import BiFNet
import torch.backends.cudnn

def getStatedict(origin, output):
    model = torch.load(origin)
    torch.save(model.state_dict(), output)


ori_weight = "./BiFNet/bifnet_starv2_48_33125_e123_92.304.pth"
dst_state_dict = "./BiFNet/bifnet_statedict_224_starv2_48_33125_e123_92.304.pth"
# model = torch.load(ori_weight)
# torch.save(model.state_dict(), dst_state_dict)

DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = BiFNet(dim=48, depth=[3, 3, 12, 5], in_chans=1, kernel_size=7, patch_size=4,
                   num_classes=16, H=224, W=224, p_h=[8, 4, 2, 1], p_w=[8, 4, 2, 1])
model.load_state_dict(torch.load(dst_state_dict))
model.to(DEVICE)
model.eval()

x1 = torch.randn(1, 1, 224, 224).to(DEVICE)
x2 = torch.randn(1, 1, 224, 224).to(DEVICE)
print(model(x1, x2).shape)